import os
import itertools
import pandas as pd
from typing import Mapping
from utils import *

pd.set_option("display.unicode.east_asian_width", True)     # print时对齐中文字符


os.makedirs(DIR_ADMISSION_SCORE_DATA, exist_ok=True)
os.makedirs(DIR_SCORE_RANK_TABLE_DATA, exist_ok=True)

def filter_columns(df, cols_map: Mapping[str, str]):
    df = df.copy()
    col_idxs = []
    # 对每一列，判断当前需要列的可能列名是否在这一列中，若在，则该列为当前需要列
    for col_name, possible_names in cols_map.items():
        candidates = df.isin(possible_names).any()
        candidates = candidates.index[candidates]     # Bool数组转为数字
        if len(candidates) == 0:
            # 将该列设置为空，并正确设置candidates
            df[col_name] = None
            # 修改这一行，创建包含列位置的Index对象
            candidates = pd.Index([df.columns.get_loc(col_name)])
            # print(f'警告：未找到列"{col_name}"，已将其设置为空')
        elif len(candidates) > 1:
            raise ValueError(f'解析失败，列"{col_name}"重复，位置{candidates}')
        col_idxs.append(candidates[0])
    # print(f'找到所需列{dict(zip(cols_map.keys(), col_idxs))}')
    res_df = df.iloc[:, col_idxs]
    res_df.columns = cols_map.keys()  # 设置列名
    return res_df

def  parse_admission_file_name(file_name):
    """解析投档线文件名，返回年份、批次和科目
    """
    year = int(file_name[:4])
    batch = None
    subject = None
    for batch_name, b in BATCH_NAMES.items():
        if batch_name in file_name:
            batch = b
            break
    for subject_name, s in SUBJECT_NAMES.items():
        if subject_name in file_name:
            subject = s
            break
    return year, batch, subject

def parse_rank_file_name(file_name):
    """解析一分一档文件名，返回年份和科目
    """
    year = int(file_name[:4])
    subject = None
    for subject_name, s in SUBJECT_NAMES.items():
        if subject_name in file_name:
            subject = s
            break
    return year, subject

def admission_file_match(file_name, year, batch, subject) -> bool:
    """判断投档线文件名是否匹配给定的年份、批次和科目
    """
    if not file_name.startswith(str(year)):
        return False
    for batch_name, b in BATCH_NAMES.items():
        if batch_name in file_name and b != batch:
            return False
    for subject_name, s in SUBJECT_NAMES.items():
        if subject_name in file_name and s != subject:
            return False
    return True

def score_rank_file_match(file_name, year, subject) -> bool:
    """判断一分一档文件名是否匹配给定的年份和科目
    """
    if not file_name.startswith(str(year)):
        return False
    for subject_name, s in SUBJECT_NAMES.items():
        if subject_name in file_name and s != subject:
            return False
    return True

years = [2024, 2023, 2022, 2021, 2020, 2019, 2018, 2017, 2016, 2015, 2014]
batches = [1, 2]
subjects = ['science', 'arts']

# 生成投档线表 
print('生成投档线表：')
admission_table = pd.DataFrame(columns=ADMISSION_COLUMNS_NAMES.keys())
for file_name in os.listdir(DIR_ADMISSION_SCORE_PAGES):
    print(f'正在处理文件 “{file_name}” ...')
    page_path = os.path.join(DIR_ADMISSION_SCORE_PAGES, file_name)
    table = pd.read_html(page_path)[0]

    # 处理投档线表，找到需要的列，将不存在的列设置为空
    filtered_table = filter_columns(table, ADMISSION_FILTER_NAMES)
    data_rows = filtered_table.iloc[:, 0].apply(lambda x: x.isdecimal())
    filtered_table = filtered_table[data_rows].iloc[:, :]
    filtered_table.reset_index(drop=True, inplace=True)
    # print(filtered_table)
    
    year, batch, subject = parse_admission_file_name(file_name)
    
    def add_row(code, name, subject, major_group, score):
        return pd.DataFrame({
            '院校代码': [code],
            '院校名称': [name],
            '年份': [year],
            '批次': [batch],
            '类别': [subject],
            '专业组': [major_group],
            '投档线': [score]
        })
    
    for row in filtered_table.itertuples(index=False):
        row_dict = dict(zip(filtered_table.columns, row))
        
        if subject is None:
            # 分别添加文史类和理工类的投档线
            if pd.notna(row_dict['文史类']):
                new_row = add_row(row_dict['院校代码'], row_dict['院校名称'], 
                                 'arts', row_dict['专业组'], row_dict['文史类'])
                admission_table = pd.concat([admission_table, new_row], ignore_index=True)
            
            if pd.notna(row_dict['理工类']):
                new_row = add_row(row_dict['院校代码'], row_dict['院校名称'], 
                                 'science', row_dict['专业组'], row_dict['理工类'])
                admission_table = pd.concat([admission_table, new_row], ignore_index=True)
        else:
            # 添加指定科目的投档线
            new_row = add_row(row_dict['院校代码'], row_dict['院校名称'], 
                             subject, row_dict['专业组'], row_dict['投档线'])
            admission_table = pd.concat([admission_table, new_row], ignore_index=True)
print(f'处理后的投档线表：\n{admission_table}')
# 所有投档线表处理完成，输出到一个总表
all_in_one_path = admission_all_in_one_path('.csv')
admission_table.to_csv(all_in_one_path, index=False)
all_in_one_path = admission_all_in_one_path('.xlsx')
admission_table.to_excel(all_in_one_path, index=False)

# 生成一分一档表
print('生成一分一档表：')
for file_name in os.listdir(DIR_SCORE_RANK_TABLE_PAGES):
    print(f'正在处理文件 “{file_name}” ...')
    page_path = os.path.join(DIR_SCORE_RANK_TABLE_PAGES, file_name)
    table = pd.read_html(page_path)[0]
    
    # 处理一分一档表
    # 找到需要的列
    table = filter_columns(table, SCORE_RANK_COLUMNS_NAMES)  
    # 找到所有第一列（分值）为数字的行，这些行是有效数据
    data_rows = table.iloc[:, 0].apply(lambda x: x.isdecimal())
    table = table[data_rows].iloc[:, :]
    table.reset_index(drop=True, inplace=True)

    # 输出表到csv及xlsx文件
    year, subject = parse_rank_file_name(file_name)
    print(f'年份: {year}, 科目: {subject}')
    data_path = score_rank_data_path(year, subject)
    table.to_csv(data_path, index=False)
    data_path = score_rank_data_path(year, subject, '.xlsx')
    table.to_excel(data_path, index=False)